# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations

import contextlib
import threading
from contextlib import nullcontext
from copy import copy
from typing import Any, Literal

import torch
from tensordict import (
    lazy_stack,
    LazyStackedTensorDict,
    MetaData,
    NonTensorStack,
    set_list_to_stack,
    TensorDict,
    TensorDictBase,
)
from tensordict.utils import _zip_strict, NestedKey
from torch import distributions as D
from torch.nn.utils.rnn import pad_sequence
from torchrl import logger as torchrl_logger
from torchrl.modules.llm.policies.common import (
    _batching,
    _extract_responses_from_full_histories,
    ChatHistory,
    LLMWrapperBase,
    LogProbs,
    Masks,
    Text,
    Tokens,
)
from torchrl.modules.utils.utils import _unpad_tensors


class TransformersWrapper(LLMWrapperBase):
    """A wrapper class for Hugging Face Transformers models, providing a consistent interface for text generation and log probability computation.

    Packing vs Padding:
        - Packing (`pad_model_input=False`):
            * More memory efficient for variable-length sequences.
            * Not all models support packed input (requires custom attention masks and position ids).
            * May be less compatible with some HuggingFace models or custom architectures.
        - Padding (`pad_model_input=True`):
            * Universally supported by all models.
            * Wastes memory for short sequences in a batch.
            * Simpler, but less efficient for highly variable-length data.
        - If unsure, use padding for maximum compatibility. Use packing for large batches of variable-length data and when your model supports it.

    Additional error handling is provided for empty and overlong sequences.

    Args:
        model (transformers.AutoModelForCausalLM | str): The Hugging Face Transformers model to wrap.
            If a string, it will be passed to `transformers.AutoModelForCausalLM.from_pretrained` (and `AutoTokenizer.from_pretrained`
            if `tokenizer` is not provided).

    Keyword Args:
        tokenizer (transformers.tokenization_utils.PreTrainedTokenizer | str | None, optional): The tokenizer to use for
            encoding and decoding text. If `None`, the tokenizer associated with the model will be used.
            If a string, it will be passed to `transformers.AutoTokenizer.from_pretrained`. Defaults to `None`.
        input_mode (str, optional): The input modality to use. Must be one of `"history"`, `"text"`, or `"tokens"`.
            Defaults to `"history"`.
        input_key (str | None, optional): The key for the input data. If `None`, defaults to
            - `("history", "prompt")` for `"history"` when `generate=True`, `("history", "full")` for `"history"` when `generate=False`
            - `("text", "prompt")` for `"text"` when `generate=True`, `("text", "full")` for `"text"` when `generate=False`
            - `("tokens", "prompt")` for `"tokens"` when `generate=True`, `("tokens", "full")` for `"tokens"` when `generate=False`
        attention_mask_key (str, optional): The key for attention masks (used in `"tokens"` mode). Defaults to `"attention_mask"`.

                    .. warning:: This argument is under development and may change in the future.

        generate (bool, optional): Whether to enable text generation. If `True`, the model will generate text based on the input.
            If `False`, only log probabilities will be computed. Defaults to `True`.
        return_log_probs (bool, optional): Whether to return log probabilities. Defaults to `False`.
        generate_kwargs (dict | None, optional): Additional arguments to pass to the model's generate method. Defaults to `None`.

            **Standardized Parameters (cross-backend compatible):**

            * **max_new_tokens** (int): Maximum number of new tokens to generate
            * **num_return_sequences** (int): Number of sequences to return
            * **temperature** (float): Sampling temperature (0.0 = deterministic, higher = more random)
            * **top_p** (float): Nucleus sampling parameter (0.0-1.0)
            * **top_k** (int): Top-k sampling parameter
            * **repetition_penalty** (float): Penalty for repeating tokens
            * **do_sample** (bool): Whether to use sampling vs greedy decoding
            * **num_beams** (int): Number of beams for beam search
            * **length_penalty** (float): Penalty for sequence length
            * **early_stopping** (bool): Whether to stop early in beam search
            * **stop_sequences** (list): Sequences that stop generation (requires custom stopping criteria)
            * **skip_special_tokens** (bool): Whether to skip special tokens in output
            * **logprobs** (bool): Whether to return log probabilities (maps to output_scores)

                .. warning:: Usage of this parameter is discouraged as it may conflict with the `generate` parameter
                    of the class.

            **Transformers-Specific Parameters:**

            * **pad_token_id** (int): Token ID for padding
            * **eos_token_id** (int): Token ID for end of sequence
            * **bad_words_ids** (list): List of token IDs to avoid
            * **force_words_ids** (list): List of token IDs to force
            * **no_repeat_ngram_size** (int): Size of n-grams to avoid repeating
            * **encoder_repetition_penalty** (float): Repetition penalty for encoder-decoder models
            * **num_beam_groups** (int): Number of beam groups for diverse beam search
            * **diversity_penalty** (float): Penalty for beam diversity
            * **output_scores** (bool): Whether to output scores
            * **return_dict_in_generate** (bool): Whether to return dict in generate

            **Legacy Parameter Support:**

            * **max_tokens** (int): Automatically converted to max_new_tokens
            * **n** (int): Automatically converted to num_return_sequences

            **Parameter Conflict Resolution:**

            When both legacy (Transformers-specific) and standardized parameter names are provided,
            a :exc:`ValueError` is raised to prevent confusion. For example:

            * If both ``max_tokens`` and ``max_new_tokens`` are passed, an error is raised
            * If both ``n`` and ``num_return_sequences`` are passed, an error is raised

            This ensures clear parameter usage and prevents unexpected behavior.

        tokenizer_kwargs (dict | None, optional): Additional arguments to pass to the tokenizer. Defaults to `None`.
        pad_output (bool, optional): Whether to pad the output sequences to a uniform length. This does not impact the underlying padding
            during call to the model. To use padding or packing during the model `forward` call, see `pad_model_input`.
            Defaults to `False`.
        pad_model_input (bool, optional): Whether to pad the model input sequences to a uniform length.
            If `False`, packing will be used instead. Packing is generally more memory efficient than padding,
            but this feature may not work with all models.
            `pad_model_input` can only be used when `generate=False`.
            This does not impact the padding of the model output - one may ask for padded output though `pad_output=True` while the model
            is called with `pad_model_input=False`.
            Defaults to `True`.
        inplace (Literal[True, False, "empty"] | None, optional): Determines how the module should handle in-place operations. Defaults to `True`.
        device (torch.device | None, optional): The device to use for computation. Defaults to `None`.
        layout (torch.layout | None, optional): The layout to use for the output tensors when `pad_output=False`. Defaults to `torch.strided`.
        num_samples (int | None, optional): The number of samples to generate. Defaults to `None` (one sample, and no batch-dimension for it).
            Can also be set via the `generate_kwargs["num_return_sequences"] = value` argument. Requires the "do_sample" argument to be set to `True` in `generate_kwargs`.
        chat_template_name (Literal["chatml_format", "qwen"] | None, optional): The name of the chat template to use when applying the chat
            template to the history. Defaults to `None`. For `input_mode="history"` only.
        chat_template (str | None, optional): The chat template to use when applying the chat template to the history.
            Defaults to `None`. For `input_mode="history"` only.
        log_probs_key (NestedKey | None, optional): The key for the log probabilities :class:`~torchrl.modules.llm.policies.LogProbs` object. Defaults to `"log_probs"`.
        text_key (NestedKey | None, optional): The key for the action :class:`~torchrl.modules.llm.policies.Text` object. Defaults to `"text"`.
        tokens_key (NestedKey | None, optional): The key for the action :class:`~torchrl.modules.llm.policies.Tokens` object. Defaults to `"tokens"`.
        masks_key (NestedKey | None, optional): The key for the action :class:`~torchrl.modules.llm.policies.Masks` object. Defaults to `"masks"`.
        history_key (NestedKey | None, optional): The key for the action :class:`~torchrl.modules.llm.policies.ChatHistory` object. Defaults to `"history"`.
        batching (bool | None, optional): Whether to enable batching. See :ref:`ref_batching` below for more details.
        min_batch_size (int | None, optional): The minimum batch size to use for batching. See :ref:`ref_batching` below for more details.
        max_batch_size (int | None, optional): The maximum batch size to use for batching. See :ref:`ref_batching` below for more details.
        batching_timeout (float, optional): The timeout for batching. See :ref:`ref_batching` below for more details.

    .. _ref_batching:
        Batching is a feature that allows the module to process multiple inputs in a single call.
        It is designed to work in a multi-threaded environment.
        To enable batching, it suffices to set `batching=True` which will set `min_batch_size` to 1 if not provided.
        If you want to set a different value for `min_batch_size` or `max_batch_size` for a fine-grained control,
        you can to set `batching=True` and then set `min_batch_size` or `max_batch_size` to a value greater or equal to 1.
        The way batching works is as follows:
        - If `min_batch_size` is not provided but `max_batch_size` is, `min_batch_size` is set to 1.
        - If `max_batch_size` is not provided but `min_batch_size` is, `max_batch_size` is set to the number of inputs in the queue.
        - When the model is called, a check is performed to see if the number of inputs in the queue is greater or equal to `min_batch_size`.
          If it is, the batch is processed immediately, while waiting for the previous batch to be processed if the model is busy.
          Otherwise, the input is added to the queue and the function waits for the batch to be completed.
          While waiting for the batch to be completed, a timeout is set to `batching_timeout` seconds such that if the batch is not
          completed after `batching_timeout` seconds, the remaining items to process are processed as is and the function returns after
          at most `batching_timeout` seconds (plus the time to finish processing the previous and current batch).

    Input Keys:
        The input key depends on both `input_mode` and `generate`:

        - If `input_mode="history"` and `generate=True`: `input_key` (defaults to `("history", "prompt")`)
        - If `input_mode="history"` and `generate=False`: `input_key` (defaults to `("history", "full")`)
        - If `input_mode="text"` and `generate=True`: `input_key` (defaults to `("text", "prompt")`)
        - If `input_mode="text"` and `generate=False`: `input_key` (defaults to `("text", "full")`)
        - If `input_mode="tokens"` and `generate=True`: `input_key` (defaults to `("tokens", "prompt")`)
        - If `input_mode="tokens"` and `generate=False`: `input_key` (defaults to `("tokens", "full")`)

    Output Keys:
        The output keys are automatically determined based on the input_mode:
        - **Tokens**: Always returned (`tokens_key`, defaults to `"tokens"`)
        - **Text**: Returned for `"text"` and `"history"` modes (`text_key`, defaults to `"text"`)
        - **History**: Returned only for `"history"` mode (`history_key`, defaults to `"history"`)
        - **Masks**: Always returned (`masks_key`, defaults to `"masks"`)
        - **Log Probs**: Returned when `return_log_probs=True` (`log_probs_key`, defaults to `"log_probs"`)

        Example output structure for `input_mode="history"`:
        ```
        TensorDict(
            text=Text(prompt=..., response=..., full=...),
            masks=Masks(all_attention_mask=..., all_assistant_mask=...),
            tokens=Tokens(prompt=..., response=..., full=...),
            log_probs=LogProbs(prompt=..., response=..., full=...),
            history=ChatHistory(prompt=..., response=..., full=...)
        )
        ```

    Example:
        >>> from transformers import AutoModelForCausalLM, AutoTokenizer
        >>> from torchrl.data.llm import History
        >>> from torchrl.modules.llm.policies import ChatHistory
        >>>
        >>> model = AutoModelForCausalLM.from_pretrained("gpt2")
        >>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
        >>>
        >>> # History input (recommended for RL environments)
        >>> wrapper = TransformersWrapper(
        ...     model,
        ...     tokenizer=tokenizer,
        ...     input_mode="history",
        ...     generate=True,
        ...     return_log_probs=True,
        ...     generate_kwargs={
        ...         "max_new_tokens": 50,  # Standardized parameter
        ...         "temperature": 0.7,
        ...         "top_p": 0.9,
        ...         "do_sample": True,
        ...     }
        ... )
        >>>
        >>> history = History.from_chats([[
        ...     {"role": "user", "content": "Hello"},
        ...     {"role": "assistant", "content": "Hi there!"}
        ... ]])
        >>> chat_history = ChatHistory(prompt=history)
        >>> result = wrapper(TensorDict(history=chat_history, batch_size=(1,)))
        >>> print(result["text"].response)  # Generated text
        >>> print(result["log_probs"].response)  # Log probabilities
        >>> print(result["history"].response)  # History with response

    Attributes:
        collector: The collector associated with the module, if it exists.

    .. seealso::
        - :class:`~torchrl.modules.llm.policies.LLMWrapperBase` (see :ref:`ref_categorical_sequential`)
        - :class:`~torchrl.modules.llm.policies.vLLMWrapper` (see :ref:`ref_vllm_wrapper`)
    """

    def __init__(
        self,
        model,
        *,
        tokenizer=None,
        input_mode: str = "history",
        input_key: str | None = None,
        attention_mask_key: str = "attention_mask",
        generate: bool = True,
        generate_kwargs: dict | None = None,
        tokenizer_kwargs: dict | None = None,
        pad_output: bool = False,
        pad_model_input: bool | None = None,
        inplace: Literal[True, False, "empty"] | None = None,
        device: torch.device | None = None,
        layout: torch.layout | None = None,
        num_samples: int | None = None,
        chat_template_name: Literal["chatml_format", "qwen"] | None = None,
        chat_template: str | None = None,
        return_log_probs: bool | None = None,
        history_key: NestedKey | None = "history",
        text_key: NestedKey | None = "text",
        tokens_key: NestedKey | None = "tokens",
        masks_key: NestedKey | None = "masks",
        log_probs_key: NestedKey | None = "log_probs",
        batching: bool | None = None,
        min_batch_size: int | None = None,
        max_batch_size: int | None = None,
        batching_timeout: float = 10.0,
    ):
        super().__init__()

        if batching and min_batch_size is None:
            min_batch_size = 1
        elif (min_batch_size is not None or max_batch_size is not None) and (
            batching is False
        ):
            raise ValueError(
                "min_batch_size and max_batch_size must be None if batching is False."
            )

        # Validate that min_batch_size <= max_batch_size when both are specified
        if min_batch_size is not None and max_batch_size is not None:
            if min_batch_size > max_batch_size:
                raise ValueError(
                    f"min_batch_size ({min_batch_size}) must be <= max_batch_size ({max_batch_size})"
                )

        self._min_batch_size = min_batch_size
        self._max_batch_size = max_batch_size
        self._batching_timeout = batching_timeout
        self._batch_queue = []
        self._futures = []
        if self.batching:
            self._batching_lock = threading.Lock()
        else:
            self._batching_lock = None

        if isinstance(model, str):
            if tokenizer is None:
                from transformers import AutoTokenizer

                tokenizer = AutoTokenizer.from_pretrained(model)

            from transformers import AutoModelForCausalLM

            model = AutoModelForCausalLM.from_pretrained(model)

        if isinstance(tokenizer, str):
            from transformers import AutoTokenizer

            tokenizer = AutoTokenizer.from_pretrained(tokenizer)

        # Validate input_mode
        if input_mode not in ["history", "text", "tokens"]:
            raise ValueError(
                f"input_mode must be one of 'history', 'text', 'tokens'. Got '{input_mode}'"
            )

        self.model = model
        self.input_mode = input_mode
        self.attention_mask_key = attention_mask_key
        self.generate = generate
        if pad_model_input is not None and generate:
            raise ValueError("pad_model_input is not supported when generate=True.")
        pad_model_input = pad_model_input if pad_model_input is not None else True
        self.pad_model_input = pad_model_input

        # Auto-determine what to return based on input mode
        self.return_history = input_mode in ("history",)
        self.return_text = input_mode in ("text", "history")
        self.return_tokens = input_mode in ("tokens", "history", "text")
        self.return_masks = True
        if return_log_probs is False and not generate:
            raise ValueError("return_log_probs must be True when generate=False.")
        return_log_probs = (
            True
            if (return_log_probs is None and generate) or (not generate)
            else bool(return_log_probs)
        )
        self.return_log_probs = return_log_probs

        self.history_key = history_key
        self.text_key = text_key
        self.tokens_key = tokens_key
        self.masks_key = masks_key
        self.log_probs_key = log_probs_key
        if not isinstance(pad_output, bool):
            raise ValueError("pad_output must be a boolean")
        self.pad_output = pad_output
        self._device = device
        if not pad_output and layout is None:
            layout = torch.strided
        self.layout = layout
        padding_value = None

        # Auto-determine input_key if not provided

        # Set input keys based on mode and generate parameter
        if input_mode == "history":
            if generate:
                self.in_keys = [
                    ("history", "prompt") if input_key is None else input_key
                ]
            else:
                self.in_keys = [("history", "full") if input_key is None else input_key]
        elif input_mode == "text":
            if generate:
                self.in_keys = [("text", "prompt") if input_key is None else input_key]
            else:
                self.in_keys = [("text", "full") if input_key is None else input_key]
        elif input_mode == "tokens":
            if generate:
                self.in_keys = [
                    ("tokens", "prompt") if input_key is None else input_key
                ]
            else:
                self.in_keys = [("tokens", "full") if input_key is None else input_key]
        self.input_key = self.in_keys[0]

        # Set output keys based on auto-determined return flags
        self.out_keys = []
        if self.return_text:
            self.out_keys.append(self.text_key)
        if self.return_masks:
            self.out_keys.append(self.masks_key)
        if self.return_tokens:
            self.out_keys.append(self.tokens_key)
        if self.return_log_probs:
            self.out_keys.append(self.log_probs_key)
        if self.return_history:
            self.out_keys.append(self.history_key)

        # Tokenizer setup
        if not tokenizer_kwargs:
            tokenizer_kwargs = {}
        else:
            tokenizer_kwargs = dict(tokenizer_kwargs)
        if not tokenizer_kwargs.setdefault("return_attention_mask", True):
            raise RuntimeError("return_attention_mask must be True")

        # We always pad, so we always return tensors
        return_tensors = "pt"
        tokenizer_kwargs.setdefault("padding", True)
        if return_tensors:
            if (
                tokenizer_kwargs.setdefault("return_tensors", return_tensors)
                != return_tensors
            ):
                raise RuntimeError

        # We always pad atm
        if tokenizer_kwargs.setdefault("padding_side", "left") != "left":
            raise RuntimeError

        self.tokenizer_kwargs = tokenizer_kwargs

        # Get tokenizer if needed
        if (
            pad_output or (input_mode in ["text", "history"] and not generate)
        ) and tokenizer is None:
            tokenizer = model.get_tokenizer()
        self.tokenizer = tokenizer

        if self.tokenizer is not None and (
            not hasattr(self.tokenizer, "pad_token") or self.tokenizer.pad_token is None
        ):
            self.tokenizer.pad_token = self.tokenizer.eos_token
        if self.tokenizer is not None:
            padding_value = self.tokenizer(self.tokenizer.pad_token)["input_ids"][0]
        self.padding_value = padding_value

        # Generate kwargs setup
        if generate_kwargs is None:
            generate_kwargs = {}
        else:
            generate_kwargs = dict(generate_kwargs)

        # Standardize common parameters
        generate_kwargs = self._standardize_generate_kwargs(generate_kwargs)

        # Extract wrapper-specific parameters
        transformers_specific_kwargs = self._get_wrapper_specific_kwargs(
            generate_kwargs, "transformers"
        )

        # Convert common parameters to Transformers format
        transformers_kwargs = {}
        for key, value in generate_kwargs.items():
            if key in self.COMMON_GENERATION_PARAMS:
                # Convert common names to Transformers names
                if key == "stop_sequences":
                    # Transformers uses stopping_criteria for stop sequences
                    # This requires custom stopping criteria implementation
                    # For now, we'll warn and skip this parameter
                    import warnings

                    warnings.warn(
                        "stop_sequences parameter is not yet fully supported in TransformersWrapper. "
                        "Use eos_token_id or implement custom stopping criteria for full support.",
                        UserWarning,
                        stacklevel=2,
                    )
                    continue
                elif key == "logprobs":
                    transformers_kwargs["output_scores"] = value
                else:
                    # Direct mapping for other common parameters
                    transformers_kwargs[key] = value

        # Add Transformers-specific parameters
        transformers_kwargs.update(transformers_specific_kwargs)

        self.num_samples = num_samples
        if (
            transformers_kwargs.get("num_return_sequences", 1) > 1
            or num_samples is not None
        ):
            if inplace in (True, "empty"):
                raise ValueError(
                    "inplace must be False (or None) when generating more than one sample."
                )
            if inplace is None:
                inplace = False
            if (
                transformers_kwargs.get("num_return_sequences", 1) > 1
                and num_samples is not None
                and transformers_kwargs.get("num_return_sequences", 1) != num_samples
            ):
                raise ValueError("num_samples differs from generate_kwargs['n'].")
            elif num_samples is None:
                self.num_samples = transformers_kwargs.get("num_return_sequences", 1)
            transformers_kwargs["num_return_sequences"] = self.num_samples
        elif inplace is None:
            inplace = True

        self.inplace = inplace

        if not generate:
            # We want only the log-probs, we generate a single token (that we then discard)
            # and retrieve the prompt log-probs
            transformers_kwargs["max_new_tokens"] = 1

        transformers_kwargs.setdefault("tokenizer", self.tokenizer)
        transformers_kwargs.setdefault("output_logits", self.return_log_probs)
        transformers_kwargs.setdefault("return_dict_in_generate", True)

        self.generate_kwargs = transformers_kwargs

        # Additional transformers-specific settings
        self.chat_template_name = chat_template_name
        self.chat_template = chat_template

        # Flag to track when we're in a get_dist call
        self._in_get_dist_call = False

    def get_new_version(self, **kwargs):
        """Returns a new version of the module with altered parameters.

        For instance, the generate parameter can be altered to enable text generation or log-probabilities computation.
        This is especially useful when one wants to avoid re-initializing the module with a new set of parameters, when the
        same parameters could be used to gather log-probs.

        Positional arguments are not supported.

        See the class constructor for more details about the parameters.
        """
        # Build the constructor arguments by using current values for missing parameters
        constructor_kwargs = {}

        # Model is always required
        constructor_kwargs["model"] = kwargs.get("model", self.model)

        # Check for each parameter and use current value if not provided
        if "tokenizer" in kwargs:
            constructor_kwargs["tokenizer"] = kwargs["tokenizer"]
        elif hasattr(self, "tokenizer"):
            constructor_kwargs["tokenizer"] = self.tokenizer

        if "input_mode" in kwargs:
            constructor_kwargs["input_mode"] = kwargs["input_mode"]
        elif hasattr(self, "input_mode"):
            constructor_kwargs["input_mode"] = self.input_mode

        if "input_key" in kwargs:
            constructor_kwargs["input_key"] = kwargs["input_key"]
        elif hasattr(self, "input_key"):
            constructor_kwargs["input_key"] = self.input_key

        if "attention_mask_key" in kwargs:
            constructor_kwargs["attention_mask_key"] = kwargs["attention_mask_key"]
        elif hasattr(self, "attention_mask_key"):
            constructor_kwargs["attention_mask_key"] = self.attention_mask_key

        if "generate" in kwargs:
            constructor_kwargs["generate"] = kwargs["generate"]
        elif hasattr(self, "generate"):
            constructor_kwargs["generate"] = self.generate

        if "generate_kwargs" in kwargs:
            constructor_kwargs["generate_kwargs"] = kwargs["generate_kwargs"]
        elif hasattr(self, "generate_kwargs"):
            constructor_kwargs["generate_kwargs"] = self.generate_kwargs

        if "pad_output" in kwargs:
            constructor_kwargs["pad_output"] = kwargs["pad_output"]
        elif hasattr(self, "pad_output"):
            constructor_kwargs["pad_output"] = self.pad_output

        if "tokenizer_kwargs" in kwargs:
            constructor_kwargs["tokenizer_kwargs"] = kwargs["tokenizer_kwargs"]
        elif hasattr(self, "tokenizer_kwargs"):
            constructor_kwargs["tokenizer_kwargs"] = self.tokenizer_kwargs
            if (
                "pad_output" in kwargs
                and kwargs.get("pad_output")
                != constructor_kwargs["tokenizer_kwargs"]["padding"]
            ):
                constructor_kwargs["tokenizer_kwargs"]["padding"] = kwargs.get(
                    "pad_output"
                )

        if "inplace" in kwargs:
            constructor_kwargs["inplace"] = kwargs["inplace"]
        elif hasattr(self, "inplace"):
            constructor_kwargs["inplace"] = self.inplace

        if "device" in kwargs:
            constructor_kwargs["device"] = kwargs["device"]
        elif hasattr(self, "_device"):
            constructor_kwargs["device"] = self._device

        if "layout" in kwargs:
            constructor_kwargs["layout"] = kwargs["layout"]
        elif hasattr(self, "layout"):
            constructor_kwargs["layout"] = self.layout

        if "num_samples" in kwargs:
            constructor_kwargs["num_samples"] = kwargs["num_samples"]
        elif hasattr(self, "num_samples"):
            constructor_kwargs["num_samples"] = self.num_samples

        if "chat_template_name" in kwargs:
            constructor_kwargs["chat_template_name"] = kwargs["chat_template_name"]
        elif hasattr(self, "chat_template_name"):
            constructor_kwargs["chat_template_name"] = self.chat_template_name

        if "chat_template" in kwargs:
            constructor_kwargs["chat_template"] = kwargs["chat_template"]
        elif hasattr(self, "chat_template"):
            constructor_kwargs["chat_template"] = self.chat_template

        if "text_key" in kwargs:
            constructor_kwargs["text_key"] = kwargs["text_key"]
        elif hasattr(self, "text_key"):
            constructor_kwargs["text_key"] = self.text_key

        if "tokens_key" in kwargs:
            constructor_kwargs["tokens_key"] = kwargs["tokens_key"]
        elif hasattr(self, "tokens_key"):
            constructor_kwargs["tokens_key"] = self.tokens_key

        if "masks_key" in kwargs:
            constructor_kwargs["masks_key"] = kwargs["masks_key"]
        elif hasattr(self, "masks_key"):
            constructor_kwargs["masks_key"] = self.masks_key

        if "log_probs_key" in kwargs:
            constructor_kwargs["log_probs_key"] = kwargs["log_probs_key"]
        elif hasattr(self, "log_probs_key"):
            constructor_kwargs["log_probs_key"] = self.log_probs_key

        # Create and return new instance
        return type(self)(**constructor_kwargs)

    @set_list_to_stack(True)
    @_batching
    def forward(
        self,
        tensordict: TensorDictBase,
        *,
        tensordict_out: TensorDictBase | None = None,
        logits_only: bool = False,
        **kwargs,
    ) -> TensorDictBase:
        tensordict_orig = tensordict
        if not tensordict.ndim:
            if tensordict_out is not None:
                raise ValueError(
                    "tensordict_out must not be provided when tensordict.ndim == 0. If this is needed, "
                    "please submit an issue on github."
                )
            # unsqueeze - squeeze the input
            return self.forward(lazy_stack([tensordict]), logits_only=logits_only)[0]
        elif tensordict.ndim > 1:
            if tensordict_out is not None:
                raise ValueError(
                    "tensordict_out must not be provided when tensordict.ndim > 1. If this is needed, "
                    "please submit an issue on github."
                )
            return self.forward(tensordict.reshape(-1), logits_only=logits_only).view(
                tensordict.shape
            )

        if not isinstance(tensordict, LazyStackedTensorDict):
            tensordict = tensordict.to_lazystack(0)

        _source_device = None
        if self._device:
            _source_device = tensordict.device
        if tensordict.device:
            tensordict = tensordict.copy().clear_device_()

        if kwargs:
            from transformers import GenerationConfig

            cfg = GenerationConfig(**kwargs)
        else:
            cfg = None

        if self.num_samples is not None:
            out = (
                TensorDict(
                    device=tensordict.device,
                    batch_size=(
                        tensordict.batch_size[0],
                        self.num_samples,
                        *tensordict.batch_size[1:],
                    ),
                )
                .to_lazystack(1)
                .to_lazystack(0)
            )
        else:
            out = TensorDict(
                device=tensordict.device, batch_size=tensordict.batch_size
            ).to_lazystack(0)

        if self.input_mode == "history":
            if self.generate:
                out = self._from_transformers_generate_history(tensordict, cfg, out)
            else:
                out = self._from_transformers_logprobs_history(
                    tensordict, cfg, out, logits_only=logits_only
                )
        elif self.input_mode == "text":
            if self.generate:
                out = self._from_transformers_generate_text(tensordict, cfg, out)
            else:
                out = self._from_transformers_logprobs_text(
                    tensordict, cfg, out, logits_only=logits_only
                )
        elif self.input_mode == "tokens":
            if self.generate:
                out = self._from_transformers_generate_tokens(tensordict, cfg, out)
            else:
                out = self._from_transformers_logprobs_tokens(
                    tensordict, cfg, out, logits_only=logits_only
                )

        if _source_device:
            out = out.to(_source_device)

        if tensordict_out is None:
            if self.inplace is True:
                # The output is the input
                tensordict_out = tensordict_orig
            elif self.inplace is False:
                # The output is the new structure
                tensordict_out = out
            elif self.inplace == "empty":
                # The output is empty
                tensordict_out = tensordict.empty()

        if tensordict_out is not None and tensordict_out is not out:
            result = tensordict_out.exclude(*self.out_keys, inplace=True)
            result.update(out, keys_to_update=self.out_keys)
        elif tensordict_out is out:
            result = out.select(*self.out_keys)
        elif self.inplace:
            result = out
            keys = list(set(self.out_keys + list(tensordict.keys(True, True))))
            result = tensordict.exclude(*self.out_keys, inplace=True).update(
                result, keys_to_update=keys
            )
        else:
            result = out
        return result

    def _from_transformers_generate_history(self, td, cfg, out) -> TensorDictBase:
        """Generate text from history input."""
        from torchrl.data.llm import History

        # Validate input
        if self.input_key not in td:
            raise ValueError(
                f"Expected '{self.input_key}' key for history input mode, "
                f"but found keys: {list(td.keys())}"
            )

        history = td.get(self.input_key)
        if not isinstance(history, History):
            raise TypeError(
                f"Expected History object for '{self.input_key}', got {type(history)}"
            )

        # Apply chat template
        tokenizer_kwargs = {}
        if self.chat_template_name is not None:
            tokenizer_kwargs.setdefault("chat_template_name", self.chat_template_name)
        if self.chat_template is not None:
            tokenizer_kwargs.setdefault("chat_template", self.chat_template)
        tokenizer_kwargs.setdefault("add_generation_prompt", True)
        text_prompt = history.apply_chat_template(
            tokenizer=self.tokenizer, **tokenizer_kwargs
        )
        if not isinstance(text_prompt, list):
            raise ValueError(
                f"Expected list of text for history input, got {type(text_prompt)}"
            )
        tokenizer_kwargs.setdefault("return_assistant_tokens_mask", False)
        tokenizer_kwargs.setdefault("tokenize", True)
        tokenizer_kwargs.setdefault("padding", False)
        tokenizer_kwargs.setdefault("return_dict", True)
        response_struct = history.apply_chat_template(
            tokenizer=self.tokenizer, **tokenizer_kwargs
        )
        tokens_prompt_padded = response_struct.get(
            "input_ids",
            as_padded_tensor=True,
            padding_value=self.padding_value,
            padding_side="left",
        )
        attention_mask_prompt_padded = response_struct.get(
            "attention_mask",
            as_padded_tensor=True,
            padding_value=0,
            padding_side="left",
        )

        if attention_mask_prompt_padded is None:
            attention_mask_prompt_padded = (
                tokens_prompt_padded != self.tokenizer.pad_token_id
            )

        result = self._generate_from_tokens(
            tokens_prompt_padded, attention_mask_prompt_padded, cfg, out
        )

        # Generate using text path
        if self.pad_output:
            result[(self.tokens_key, "prompt")] = (
                tokens_prompt_padded
                if not self.num_samples
                else tokens_prompt_padded.unsqueeze(1).repeat(1, self.num_samples, 1)
            )
        else:
            tokens_prompt_unpadded = response_struct.get(
                "input_ids",
                as_nested_tensor=True,
            )
            if not self.num_samples:
                result[(self.tokens_key, "prompt")] = tokens_prompt_unpadded
            else:
                for r in result.unbind(1):
                    r[(self.tokens_key, "prompt")] = tokens_prompt_unpadded

        text_result = Text._from_tensordict(result.empty())
        result.set(self.text_key, text_result)
        if not self.num_samples:
            text_result.prompt = text_prompt
        else:
            for r in result.unbind(1):
                r[self.text_key, "prompt"] = text_prompt
        with result.view(-1) as result_flat:
            if self.pad_output:
                tokens_full_padded = result_flat.get(
                    (self.tokens_key, "full"),
                    as_padded_tensor=True,
                    padding_side="right",
                    padding_value=self.padding_value,
                )
                if tokens_full_padded is None:
                    raise ValueError("tokens_full_padded is None")
                text_full = self.tokenizer.batch_decode(
                    tokens_full_padded, skip_special_tokens=False
                )
            else:
                tokens_full_unpadded = result_flat.get(
                    (self.tokens_key, "full"), as_list=True
                )
                if tokens_full_unpadded is None:
                    raise ValueError("tokens_full_unpadded is None")
                text_full = self.tokenizer.batch_decode(
                    tokens_full_unpadded, skip_special_tokens=False
                )
            text_prompt = result_flat[self.text_key, "prompt"]
            text_response = [
                txt[len(prompt) :]
                for txt, prompt in _zip_strict(text_full, text_prompt)
            ]
            result_flat.set((self.text_key, "full"), text_full)
            result_flat.set((self.text_key, "response"), text_response)
        # Now parse the full text back to a history object, and use the extra history objects
        # as response
        history_chat = ChatHistory._from_tensordict(result.empty())
        if self.num_samples is None:
            history_chat.prompt = history
        else:
            for h in history_chat.unbind(1):
                h.prompt = history
        with history_chat.view(-1) as history_chat_flat:
            prompt_histories = history_chat_flat.prompt
            # Extract response histories from full text
            h_responses = _extract_responses_from_full_histories(
                text_full, prompt_histories, self.chat_template_name, self.tokenizer
            )
            history_chat_flat.response = h_responses
        result.set(self.history_key, history_chat)
        return result

    def _from_transformers_logprobs_history(self, td, cfg, out, logits_only=False):
        """Compute log-probs from history input."""
        from torchrl.data.llm import History

        # Validate input
        if self.input_key not in td:
            raise ValueError(
                f"Expected '{self.input_key}' key for history input mode, "
                f"but found keys: {list(td.keys())}"
            )

        history = td.get(self.input_key)
        if not isinstance(history, History):
            raise TypeError(
                f"Expected History object for '{self.input_key}', got {type(history)}"
            )

        # Apply chat template
        tokenizer_kwargs = {}
        if self.chat_template_name is not None:
            tokenizer_kwargs.setdefault("chat_template_name", self.chat_template_name)
        if self.chat_template is not None:
            tokenizer_kwargs.setdefault("chat_template", self.chat_template)
        tokenizer_kwargs.setdefault("add_generation_prompt", False)
        text_full = history.apply_chat_template(
            tokenizer=self.tokenizer, **tokenizer_kwargs
        )

        tokenizer_kwargs.setdefault("return_assistant_tokens_mask", True)
        tokenizer_kwargs.setdefault("tokenize", True)
        tokenizer_kwargs.setdefault("padding", False)
        tokenizer_kwargs.setdefault("return_dict", True)

        with torch.device(self._device) if self._device is not None else nullcontext():
            response_tokens = history.apply_chat_template(
                tokenizer=self.tokenizer, **tokenizer_kwargs
            )
        if not isinstance(response_tokens, TensorDictBase):
            raise ValueError(
                f"Expected TensorDictBase for history input, got {type(response_tokens)}"
            )
        result = self._logprobs_from_history_tokens(
            response_tokens, cfg, out, logits_only=logits_only
        )
        text_result = Text._from_tensordict(result.empty())
        result.set(self.text_key, text_result)
        result[self.text_key, "full"] = text_full
        result.set(self.history_key, ChatHistory(full=history))
        return result

    def _cat_text(self, text, response_text):
        """Concatenate text and response text."""
        if isinstance(text, list):
            return [self._cat_text(t, t_) for t, t_ in _zip_strict(text, response_text)]
        else:
            return text + response_text

    def _generate_from_text(self, text, cfg, out) -> TensorDictBase:
        """Generate text from text input."""
        pad_val = self.tokenizer.pad_token_id

        # Convert text to list format
        if isinstance(text, str):
            text = [text]
        elif not isinstance(text, list):
            text = text.tolist()

        tokenizer_kwargs = dict(self.tokenizer_kwargs)
        tokenizer_kwargs.setdefault("padding", True)

        with torch.device(
            self._device
        ) if self._device is not None else contextlib.nullcontext():
            tokens_in = self.tokenizer(text, **tokenizer_kwargs)
        if self._device is not None:
            tokens_in = tokens_in.to(self._device)
        # We are going to map this tokens_in to a tensordict to facilitate the padding in case we need it
        tokens_in = dict(tokens_in)
        for k, v in dict(tokens_in).items():
            if isinstance(v, list):
                if isinstance(v[0], torch.Tensor):
                    v = torch.nested.nested_tensor(v)
                else:
                    v = torch.nested.nested_tensor([torch.tensor(t) for t in v])
            tokens_in[k] = v
        tokens_in = (
            TensorDict(batch_size=tokens_in["input_ids"].size(0))
            .to_lazystack(0)
            .update(tokens_in)
        )
        tokens_prompt_padded = tokens_in.get(
            "input_ids",
            as_padded_tensor=True,
            padding_side="left",
            padding_value=pad_val,
        )
        attention_mask_prompt_padded = tokens_in.get(
            "attention_mask",
            as_padded_tensor=True,
            padding_side="left",
            padding_value=0,
        )

        if cfg is not None:
            kwargs = copy(self.generate_kwargs)
            kwargs["generation_config"] = cfg
        else:
            kwargs = self.generate_kwargs

        tokens_out = self.model.generate(
            input_ids=tokens_prompt_padded,
            attention_mask=attention_mask_prompt_padded,
            **kwargs,
        )
        tokens_full_padded = tokens_out["sequences"]
        tokens_response_padded = tokens_full_padded[
            ..., tokens_prompt_padded.shape[-1] :
        ]

        attention_mask_response_padded = tokens_response_padded != pad_val
        if self.num_samples:
            attention_mask_full_padded = torch.cat(
                [
                    attention_mask_prompt_padded.repeat_interleave(
                        self.num_samples, dim=0
                    ),
                    attention_mask_response_padded,
                ],
                dim=-1,
            )
        else:
            attention_mask_full_padded = torch.cat(
                [attention_mask_prompt_padded, attention_mask_response_padded], dim=-1
            )
        tokens_response_unpadded = _unpad_tensors(
            tokens_response_padded, attention_mask_response_padded, as_nested=False
        )

        if self.return_log_probs:
            # These are only for the new tokens, not for the prompt - to get that, we'd need to run the forward pass again
            logits = torch.stack(list(tokens_out["logits"]), 1)
            log_probs, logits = self._log_probs_generate(
                tokens_response_padded, logits, pad_val=-100, pad=False
            )

        response_text = self.tokenizer.batch_decode(
            tokens_response_unpadded, skip_special_tokens=False
        )

        # Build output TensorClass objects
        if self.num_samples is not None:
            text = [txt for txt in text for _ in range(self.num_samples)]
        text_obj = Text._from_tensordict(out.empty())
        with text_obj.view(-1) as text_obj_flat:
            text_obj_flat.prompt = text
            text_obj_flat.response = response_text
            text_obj_flat.full = self._cat_text(text, response_text)
        out.set(self.text_key, text_obj)

        tokens_obj = Tokens._from_tensordict(out.empty())
        if self.pad_output:
            prompt = tokens_prompt_padded
        else:
            prompt = _unpad_tensors(
                tokens_prompt_padded, attention_mask_prompt_padded, as_nested=False
            )
        if tokens_obj.ndim == 2:
            for i in range(self.num_samples):
                tokens_obj[:, i].prompt = prompt
        else:
            tokens_obj.prompt = prompt
        with tokens_obj.view(-1) as tokens_obj_flat:
            if not self.pad_output:
                tokens_obj_flat.response = tokens_response_unpadded
                tokens_full_unpadded = _unpad_tensors(
                    tokens_full_padded, attention_mask_full_padded, as_nested=False
                )
                tokens_obj_flat.full = tokens_full_unpadded
            else:
                tokens_obj_flat.response = tokens_response_padded
                tokens_obj_flat.full = tokens_full_padded
        tokens_obj.padded = MetaData(self.pad_output)
        out.set(self.tokens_key, tokens_obj)

        masks_obj = Masks._from_tensordict(out.empty())
        if out.ndim == 2:
            attention_mask_full_padded = attention_mask_full_padded.unflatten(
                0, (-1, self.num_samples)
            )
        if self.pad_output:
            masks_obj.all_attention_mask = attention_mask_full_padded.bool()
        else:
            if out.ndim == 2:
                with tokens_obj.view(-1) as tokens_obj_flat, masks_obj.view(
                    -1
                ) as masks_obj_flat:
                    attention_mask_full_unpadded = attention_mask_full_padded.flatten(
                        0, 1
                    )
                    attention_mask_full_unpadded = _unpad_tensors(
                        attention_mask_full_unpadded.bool(),
                        attention_mask_full_padded.flatten(0, 1),
                        as_nested=False,
                    )
                    masks_obj_flat.all_attention_mask = attention_mask_full_unpadded
            else:
                attention_mask_full_unpadded = _unpad_tensors(
                    attention_mask_full_padded.bool(),
                    attention_mask_full_padded,
                    as_nested=False,
                )
                masks_obj.all_attention_mask = attention_mask_full_unpadded
        masks_obj.all_assistant_mask = None
        masks_obj.padded = MetaData(self.pad_output)
        out.set(self.masks_key, masks_obj)

        if self.return_log_probs:
            log_probs_obj = LogProbs._from_tensordict(out.empty())
            with log_probs_obj.view(-1) as log_probs_obj_flat:
                # Unfortunate but we only have the log-probs for the new tokens, not for the prompt - to get that, we'd need to run the forward pass again
                if self.pad_output:
                    log_probs_obj_flat.prompt = None
                    log_probs_obj_flat.response = log_probs
                    log_probs_obj_flat.full = None
                else:
                    log_probs_unpadded = _unpad_tensors(
                        log_probs, attention_mask_response_padded, as_nested=False
                    )
                    log_probs_obj_flat.prompt = None
                    log_probs_obj_flat.response = log_probs_unpadded
                    log_probs_obj_flat.full = None
            log_probs_obj.padded = MetaData(self.pad_output)
            out.set(self.log_probs_key, log_probs_obj)

        # Add logits to output if we're in a get_dist call
        if self._in_get_dist_call:
            if self.pad_output:
                out.set("logits", logits)
            else:
                logits_full_unpadded = _unpad_tensors(
                    logits, attention_mask_full_padded, as_nested=False
                )
                out.set("logits", logits_full_unpadded)

        return out

    def _cat_tensors(
        self,
        tokens: torch.Tensor | list[torch.Tensor],
        response_tokens: torch.Tensor | list[torch.Tensor],
        cast: torch.dtype | None = None,
    ):
        """Concatenate tokens and response tokens."""
        if isinstance(tokens, list) or isinstance(response_tokens, list):
            return [
                self._cat_tensors(t, t_, cast=cast)
                for t, t_ in _zip_strict(tokens, response_tokens)
            ]
        else:
            result = torch.cat([tokens, response_tokens], dim=-1)
            if cast is not None:
                result = result.to(cast)
            return result

    def _logprobs_from_history_tokens(
        self, response_tokens, cfg, out, logits_only=False
    ):
        """Compute log-probs from history tokens."""
        pad_val = self.tokenizer.pad_token_id

        if cfg is not None:
            kwargs = copy(self.generate_kwargs)
            kwargs["generation_config"] = cfg
        else:
            kwargs = self.generate_kwargs

        # non-packed forward pass
        if self.pad_model_input:
            # unfortunately HF wants us to use padded tensors
            tokens_full_padded = response_tokens.get(
                "input_ids",
                as_padded_tensor=True,
                padding_side="left",
                padding_value=pad_val,
            )
            if not isinstance(tokens_full_padded, torch.Tensor):
                raise ValueError(
                    f"Expected Tensor for tokens_full_padded, got {type(tokens_full_padded)}"
                )
            attention_mask_full_padded = response_tokens.get(
                "attention_mask",
                as_padded_tensor=True,
                padding_side="left",
                padding_value=0,
            )
            if not isinstance(attention_mask_full_padded, torch.Tensor):
                raise ValueError(
                    f"Expected Tensor for attention_mask_full_padded, got {type(attention_mask_full_padded)}"
                )

            (
                log_probs_full_padded,
                logits_full_padded,
            ) = self._model_forward_with_padded_sequences(
                tokens_full_padded,
                attention_mask_full_padded,
                pad_val=pad_val,
                logits_only=logits_only,
                **kwargs,
            )
        else:
            # unfortunately HF wants us to use padded tensors
            tokens_full_unpadded = response_tokens.get(
                "input_ids",
                as_nested_tensor=True,
                layout=torch.jagged,
            )
            attention_mask_full_unpadded = response_tokens.get(
                "attention_mask",
                as_nested_tensor=True,
                layout=torch.jagged,
            )
            (
                log_probs_full_unpadded,
                logits_full_unpadded,
            ) = self._model_forward_with_packed_sequences(
                # TODO: no padding if we don't need to
                tokens_full_unpadded,
                attention_mask_full_unpadded,
                pad=False,
                logits_only=logits_only,
                **kwargs,
            )
            tokens_full_padded = pad_sequence(
                tokens_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=pad_val,
                padding_side="left",
            )
            attention_mask_full_padded = pad_sequence(
                attention_mask_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=0,
                padding_side="left",
            )
            if log_probs_full_unpadded is not None:
                log_probs_full_padded = pad_sequence(
                    log_probs_full_unpadded.unbind(0),
                    batch_first=True,
                    padding_value=0.0,
                    padding_side="left",
                )
            else:
                log_probs_full_padded = None
            logits_full_padded = pad_sequence(
                logits_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=0.0,
                padding_side="left",
            )
        # Build output TensorClass objects
        text_obj = Text._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        text_obj.prompt = None
        text_obj.response = None
        text_obj.full = None
        out.set(self.text_key, text_obj)

        all_assistant_mask_padded = response_tokens.get(
            "assistant_masks",
            as_padded_tensor=True,
            padding_side="left",
            padding_value=0,
        )
        if all_assistant_mask_padded is not None:
            all_assistant_mask_padded = all_assistant_mask_padded.bool()
        masks_obj = Masks._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        if self.pad_output:
            masks_obj.all_attention_mask = attention_mask_full_padded.bool()
            if all_assistant_mask_padded is not None:
                masks_obj.all_assistant_mask = all_assistant_mask_padded
        else:
            masks_obj.all_attention_mask = _unpad_tensors(
                attention_mask_full_padded.bool(),
                attention_mask_full_padded,
                as_nested=False,
            )
            if all_assistant_mask_padded is not None:
                masks_obj.all_assistant_mask = _unpad_tensors(
                    all_assistant_mask_padded,
                    attention_mask_full_padded,
                    as_nested=False,
                )
        masks_obj.padded = MetaData(self.pad_output)
        out.set(self.masks_key, masks_obj)

        tokens_obj = Tokens._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        if self.pad_output:
            tokens_obj.full = tokens_full_padded
        else:
            input_ids_full_unpadded = _unpad_tensors(
                tokens_full_padded, attention_mask_full_padded, as_nested=False
            )
            tokens_obj.full = input_ids_full_unpadded
        tokens_obj.response = None
        tokens_obj.padded = MetaData(self.pad_output)
        out.set(self.tokens_key, tokens_obj)

        if not logits_only:
            log_probs_obj = LogProbs._from_tensordict(
                TensorDict(batch_size=out.batch_size).to_lazystack(0)
            )
            if self.pad_output:
                log_probs_obj.full = log_probs_full_padded
            else:
                log_probs_full_unpadded = _unpad_tensors(
                    log_probs_full_padded, attention_mask_full_padded, as_nested=False
                )
                log_probs_obj.full = log_probs_full_unpadded
            log_probs_obj.response = None
            log_probs_obj.padded = MetaData(self.pad_output)
            out.set(self.log_probs_key, log_probs_obj)

        # Add logits to output if we're in a get_dist call
        if self._in_get_dist_call:
            if self.pad_output:
                out.set("logits", logits_full_padded)
            else:
                logits_full_unpadded = _unpad_tensors(
                    logits_full_padded, attention_mask_full_padded, as_nested=False
                )
                out.set("logits", logits_full_unpadded)

        return out

    def _from_transformers_generate_text(self, td, cfg, out) -> TensorDictBase:
        """Generate text from text input."""
        # Validate input
        if self.input_key not in td:
            raise ValueError(
                f"Expected '{self.input_key}' key for text input mode, "
                f"but found keys: {list(td.keys())}"
            )

        text = td.get(self.input_key)
        if text is None:
            raise ValueError(f"Expected '{self.input_key}' key for text input mode")
        if isinstance(text, NonTensorStack):
            text = text.tolist()
        if not isinstance(text, list):
            raise ValueError(f"Expected list of text for text input, got {type(text)}")
        return self._generate_from_text(text, cfg, out)

    def _from_transformers_logprobs_text(self, td, cfg, out, logits_only=False):
        """Compute log-probs from text input."""
        # Validate input
        if self.input_key not in td:
            raise ValueError(
                f"Expected '{self.input_key}' key for text input mode, "
                f"but found keys: {list(td.keys())}"
            )

        text = td.get(self.input_key)
        if isinstance(text, NonTensorStack):
            text = text.tolist()
        if text is None:
            raise ValueError(f"Expected '{self.input_key}' key for text input mode")
        if not isinstance(text, list):
            raise ValueError(f"Expected list of text for text input, got {type(text)}")
        # Tokenize the text
        if self.tokenizer is None:
            raise ValueError(
                "Tokenizer is required for log-probs computation with text input"
            )

        # Convert text to list format
        if isinstance(text, str):
            text = [text]
        elif not isinstance(text, list):
            text = text.tolist()

        # Tokenize the text
        tokenizer_kwargs = dict(self.tokenizer_kwargs)
        with torch.device(
            self._device
        ) if self._device is not None else contextlib.nullcontext():
            tokens_in = self.tokenizer(text, **tokenizer_kwargs)

        if cfg is not None:
            kwargs = copy(self.generate_kwargs)
            kwargs["generation_config"] = cfg
        else:
            kwargs = self.generate_kwargs

        # We are going to map this tokens_in to a tensordict to facilitate the padding in case we need it
        tokens_in = (
            TensorDict(batch_size=len(tokens_in["input_ids"]))
            .to_lazystack(0)
            .update(dict(tokens_in))
        )
        pad_val = self.padding_value

        if self.pad_model_input:
            tokens_full_padded = tokens_in.get(
                "input_ids",
                as_padded_tensor=True,
                padding_side="left",
                padding_value=pad_val,
            )
            attention_mask_full_padded = tokens_in.get(
                "attention_mask",
                as_padded_tensor=True,
                padding_side="left",
                padding_value=0,
            )

            (
                log_probs_full_padded,
                logits_full_padded,
            ) = self._model_forward_with_padded_sequences(
                tokens_full_padded,
                attention_mask_full_padded,
                pad_val=pad_val,
                logits_only=logits_only,
                **kwargs,
            )
        else:
            # packed forward pass
            tokens_full_unpadded = tokens_in.get(
                "input_ids",
                as_nested_tensor=True,
                layout=torch.jagged,
            )
            attention_mask_full_unpadded = tokens_in.get(
                "attention_mask",
                as_nested_tensor=True,
                layout=torch.jagged,
            )
            (
                log_probs_full_unpadded,
                logits_full_unpadded,
            ) = self._model_forward_with_packed_sequences(
                tokens_full_unpadded, attention_mask_full_unpadded, pad=False, **kwargs
            )
            tokens_full_padded = pad_sequence(
                tokens_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=pad_val,
                padding_side="left",
            )
            attention_mask_full_padded = pad_sequence(
                attention_mask_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=0,
                padding_side="left",
            )
            log_probs_full_padded = pad_sequence(
                log_probs_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=0.0,
                padding_side="left",
            )
            logits_full_padded = pad_sequence(
                logits_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=0.0,
                padding_side="left",
            )

        # Build output TensorClass objects
        text_obj = Text._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        text_obj.prompt = None
        text_obj.response = None
        text_obj.full = text
        out.set(self.text_key, text_obj)

        tokens_obj = Tokens._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        if self.pad_output:
            tokens_obj.full = tokens_full_padded
        else:
            input_ids_full_unpadded = _unpad_tensors(
                tokens_full_padded, attention_mask_full_padded, as_nested=False
            )
            tokens_obj.full = input_ids_full_unpadded
        tokens_obj.response = None
        tokens_obj.padded = MetaData(self.pad_output)
        out.set(self.tokens_key, tokens_obj)

        masks_obj = Masks._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        if self.pad_output:
            masks_obj.all_attention_mask = attention_mask_full_padded.bool()
            masks_obj.all_assistant_mask = td.get(("masks", "all_assistant_mask"))
        else:
            attention_mask_full_unpadded = _unpad_tensors(
                attention_mask_full_padded.bool(),
                attention_mask_full_padded,
                as_nested=False,
            )
            masks_obj.all_attention_mask = attention_mask_full_unpadded
            masks_obj.all_assistant_mask = td.get(
                ("masks", "all_assistant_mask"), as_list=True
            )
        masks_obj.padded = MetaData(self.pad_output)
        out.set(self.masks_key, masks_obj)

        if not logits_only:
            log_probs_obj = LogProbs._from_tensordict(
                TensorDict(batch_size=out.batch_size).to_lazystack(0)
            )
            if self.pad_output:
                log_probs_obj.full = log_probs_full_padded
            else:
                log_probs_full_unpadded = _unpad_tensors(
                    log_probs_full_padded, attention_mask_full_padded, as_nested=False
                )
                log_probs_obj.full = log_probs_full_unpadded
            log_probs_obj.response = None
            log_probs_obj.padded = MetaData(self.pad_output)
            out.set(self.log_probs_key, log_probs_obj)

        # Add logits to output if we're in a get_dist call
        if self._in_get_dist_call:
            if self.pad_output:
                out.set("logits", logits_full_padded)
            else:
                logits_full_unpadded = _unpad_tensors(
                    logits_full_padded, attention_mask_full_padded, as_nested=False
                )
                out.set("logits", logits_full_unpadded)

        return out

    def _from_transformers_generate_tokens(
        self, td: TensorDictBase, cfg: dict | None, out: TensorDictBase
    ) -> TensorDictBase:
        """Generate text from tokens input."""
        # Validate input
        if self.input_key not in td:
            raise ValueError(
                f"Expected '{self.input_key}' key for tokens input mode, "
                f"but found keys: {list(td.keys())}"
            )

        pad_val = self.tokenizer.pad_token_id

        input_ids_prompt_padded = td.get(
            self.input_key,
            as_padded_tensor=True,
            padding_side="left",
            padding_value=pad_val,
        )
        attention_mask_prompt_padded = td.get(
            ("masks", "all_attention_mask"),
            as_padded_tensor=True,
            padding_side="left",
            padding_value=False,
        )
        if attention_mask_prompt_padded is None:
            attention_mask_prompt_padded = td.get(
                self.attention_mask_key,
                as_padded_tensor=True,
                padding_side="left",
                padding_value=False,
            )
            if attention_mask_prompt_padded is None:
                attention_mask_prompt_padded = input_ids_prompt_padded != pad_val
        return self._generate_from_tokens(
            input_ids_prompt_padded, attention_mask_prompt_padded, cfg, out
        )

    def _generate_from_tokens(
        self,
        tokens_prompt_padded: torch.Tensor,
        attention_mask_prompt_padded: torch.Tensor,
        cfg: dict | None,
        out: TensorDictBase,
    ) -> TensorDictBase:
        if cfg is not None:
            kwargs = copy(self.generate_kwargs)
            kwargs["generation_config"] = cfg
        else:
            kwargs = self.generate_kwargs

        tokens_out_struct = self.model.generate(
            input_ids=tokens_prompt_padded,
            attention_mask=attention_mask_prompt_padded,
            **kwargs,
        )
        tokens_full_padded = tokens_out_struct["sequences"]
        tokens_response_padded = tokens_full_padded[:, tokens_prompt_padded.shape[-1] :]
        pad_val = getattr(self.tokenizer, "pad_token_id", None)
        if pad_val is None:
            pad_val = self.padding_value
        attention_mask_reponse_padded = tokens_response_padded != pad_val
        attention_mask_full_padded = tokens_full_padded != pad_val
        tokens_response_unpadded = _unpad_tensors(
            tokens_response_padded, attention_mask_reponse_padded, as_nested=False
        )

        if self.return_log_probs:
            # These are only for the new tokens, not for the prompt - to get that, we'd need to run the forward pass again
            logits_response_padded = tokens_out_struct["logits"]
            logits_response_padded = torch.stack(list(logits_response_padded), 1)
            (
                log_probs_response_padded,
                logits_response_padded,
            ) = self._log_probs_generate(
                tokens_response_padded,
                logits_response_padded,
                pad_val=pad_val,
                pad=False,
            )

        response_text = self.tokenizer.batch_decode(
            tokens_response_unpadded, skip_special_tokens=False
        )

        # Build output TensorClass objects
        text_obj = Text._from_tensordict(out.empty())
        text_obj.prompt = None  # We don't have text in tokens mode
        with text_obj.view(-1) as text_obj_flat:
            text_obj_flat.response = response_text
        text_obj.full = None  # we don't have text in tokens mode so no all_text either
        out.set(self.text_key, text_obj)

        tokens_obj = Tokens._from_tensordict(out.empty())
        if not self.pad_output:
            input_ids_prompt_unpadded = _unpad_tensors(
                tokens_prompt_padded,
                attention_mask_prompt_padded,
                as_nested=False,
            )
        if self.num_samples is not None:
            # replicate tokens
            for i in range(self.num_samples):
                tokens_obj[:, i].prompt = (
                    input_ids_prompt_unpadded
                    if not self.pad_output
                    else tokens_prompt_padded
                )
        else:
            tokens_obj.prompt = (
                input_ids_prompt_unpadded
                if not self.pad_output
                else tokens_prompt_padded
            )
        with tokens_obj.view(-1) as tokens_obj_flat:
            if self.pad_output:
                tokens_obj_flat.response = tokens_response_padded
                tokens_obj_flat.full = tokens_full_padded
            else:
                tokens_obj_flat.response = tokens_response_unpadded
                tokens_full_unpadded = _unpad_tensors(
                    tokens_full_padded, attention_mask_full_padded, as_nested=False
                )
                tokens_obj_flat.full = tokens_full_unpadded
        tokens_obj.padded = MetaData(self.pad_output)
        out.set(self.tokens_key, tokens_obj)

        masks_obj = Masks._from_tensordict(out.empty())
        if out.ndim == 2:
            attention_mask_full_padded = attention_mask_full_padded.unflatten(
                0, (-1, self.num_samples)
            )
        if self.pad_output:
            # Get "real" attention masks
            masks_obj.all_attention_mask = attention_mask_full_padded.bool()
        else:
            # Get "real" attention masks
            # We can use select to avoid batch-size problems
            _td = torch.ones_like(
                out.select(("tokens", "full"))
                .copy()
                .rename_key_(("tokens", "full"), "all_attention_mask")
            ).bool()
            del _td["tokens"]
            masks_obj.update(_td)
        masks_obj.all_assistant_mask = None
        masks_obj.padded = MetaData(self.pad_output)
        out.set(self.masks_key, masks_obj)

        if self.return_log_probs:
            log_probs_obj = LogProbs._from_tensordict(out.empty())
            if self.num_samples is None:
                if self.pad_output:
                    log_probs_obj.response = log_probs_response_padded
                else:
                    log_probs_response_unpadded = _unpad_tensors(
                        log_probs_response_padded,
                        attention_mask_reponse_padded,
                        as_nested=False,
                    )
                    log_probs_obj.response = log_probs_response_unpadded
            else:
                with log_probs_obj.view(-1) as log_probs_obj_flat:
                    if self.pad_output:
                        log_probs_obj_flat.response = log_probs_response_padded
                    else:
                        log_probs_response_unpadded = _unpad_tensors(
                            log_probs_response_padded,
                            attention_mask_reponse_padded,
                            as_nested=False,
                        )
                        log_probs_obj_flat.response = log_probs_response_unpadded
            log_probs_obj.padded = MetaData(self.pad_output)
            out.set(self.log_probs_key, log_probs_obj)

        return out

    def _from_transformers_logprobs_tokens(
        self,
        td: TensorDictBase,
        cfg: dict | None,
        out: TensorDictBase,
        logits_only=False,
    ) -> TensorDictBase:
        """Compute log-probs from tokens input."""
        # Validate input
        if self.input_key not in td:
            raise ValueError(
                f"Expected '{self.input_key}' key for tokens input mode, "
                f"but found keys: {list(td.keys(isinstance(self.input_key, tuple)))}"
            )

        pad_val = self.tokenizer.pad_token_id

        if cfg is not None:
            kwargs = copy(self.generate_kwargs)
            kwargs["generation_config"] = cfg
        else:
            kwargs = self.generate_kwargs

        if self.pad_model_input:
            tokens_full_padded = td.get(
                self.input_key,
                as_padded_tensor=True,
                padding_side="left",
                padding_value=pad_val,
            )
            # Attention mask: try first the regular entry, then the key provided in the constructor, finally fallback on eager attention mask
            attention_mask_full_padded = td.get(
                ("masks", "all_attention_mask"),
                as_padded_tensor=True,
                padding_side="left",
                padding_value=False,
            )
            if attention_mask_full_padded is None:
                attention_mask_full_padded = td.get(
                    self.attention_mask_key,
                    as_padded_tensor=True,
                    padding_side="left",
                    padding_value=False,
                )
                if attention_mask_full_padded is None:
                    attention_mask_full_padded = tokens_full_padded != pad_val

            (
                log_probs_full_padded,
                logits_full_padded,
            ) = self._model_forward_with_padded_sequences(
                tokens_full_padded,
                attention_mask_full_padded,
                pad_val=pad_val,
                logits_only=logits_only,
                **kwargs,
            )
        else:
            # packed forward pass
            # unfortunately HF wants us to use padded tensors
            tokens_full_unpadded = td.get(
                self.input_key,
                as_nested_tensor=True,
                layout=torch.jagged,
            )
            if tokens_full_unpadded is None:
                raise ValueError(
                    f"Expected '{self.input_key}' key for tokens input mode, but found keys: {list(td.keys())}"
                )
            # Attention mask: try first the regular entry, then the key provided in the constructor, finally fallback on eager attention mask
            attention_mask_full_unpadded = td.get(
                ("masks", "all_attention_mask"),
                as_nested_tensor=True,
                layout=torch.jagged,
            )
            if attention_mask_full_unpadded is None:
                attention_mask_full_unpadded = td.get(
                    self.attention_mask_key,
                    as_nested_tensor=True,
                    layout=torch.jagged,
                )
                if attention_mask_full_unpadded is None:
                    # does this even work?
                    attention_mask_full_unpadded = tokens_full_unpadded != pad_val

            (
                log_probs_full_unpadded,
                logits_full_unpadded,
            ) = self._model_forward_with_packed_sequences(
                # TODO: no padding if we don't need to
                tokens_full_unpadded,
                attention_mask_full_unpadded,
                pad=False,
                logits_only=logits_only,
                **kwargs,
            )
            tokens_full_padded = pad_sequence(
                tokens_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=pad_val,
                padding_side="left",
            )
            attention_mask_full_padded = pad_sequence(
                attention_mask_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=0,
                padding_side="left",
            )
            if log_probs_full_unpadded is not None:
                log_probs_full_padded = pad_sequence(
                    log_probs_full_unpadded.unbind(0),
                    batch_first=True,
                    padding_value=0.0,
                    padding_side="left",
                )
            else:
                log_probs_full_padded = None
            logits_full_padded = pad_sequence(
                logits_full_unpadded.unbind(0),
                batch_first=True,
                padding_value=0.0,
                padding_side="left",
            )

        # Build output TensorClass objects
        text_obj = Text._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        text_obj.prompt = None
        text_obj.response = None
        text_obj.full = None
        out.set(self.text_key, text_obj)

        tokens_obj = Tokens._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        if not self.pad_output:
            input_ids_full_unpadded = _unpad_tensors(
                tokens_full_padded, attention_mask_full_padded, as_nested=False
            )
            tokens_obj.full = input_ids_full_unpadded
        else:
            tokens_obj.full = tokens_full_padded
        tokens_obj.response = None
        tokens_obj.padded = MetaData(self.pad_output)
        out.set(self.tokens_key, tokens_obj)

        masks_obj = Masks._from_tensordict(
            TensorDict(batch_size=out.batch_size).to_lazystack(0)
        )
        if self.pad_output:
            masks_obj.all_attention_mask = attention_mask_full_padded.bool()
            masks_obj.all_assistant_mask = td.get(("masks", "all_assistant_mask"))
        else:
            masks_obj.all_attention_mask = _unpad_tensors(
                attention_mask_full_padded.bool(),
                attention_mask_full_padded,
                as_nested=False,
            )
            masks_obj.all_assistant_mask = td.get(
                ("masks", "all_assistant_mask"), as_list=True
            )

        masks_obj.padded = MetaData(self.pad_output)
        out.set(self.masks_key, masks_obj)

        if not logits_only:
            log_probs_obj = LogProbs._from_tensordict(
                TensorDict(batch_size=out.batch_size).to_lazystack(0)
            )
            if self.pad_output:
                log_probs_obj.full = log_probs_full_padded
            else:
                log_probs_full_unpadded = _unpad_tensors(
                    log_probs_full_padded, attention_mask_full_padded, as_nested=False
                )
                log_probs_obj.full = log_probs_full_unpadded
            log_probs_obj.response = None
            log_probs_obj.padded = MetaData(self.pad_output)
            out.set(self.log_probs_key, log_probs_obj)

        # Add logits to output if we're in a get_dist call
        if self._in_get_dist_call:
            if self.pad_output:
                out.set("logits", logits_full_padded)
            else:
                logits_full_unpadded = _unpad_tensors(
                    logits_full_padded, attention_mask_full_padded, as_nested=False
                )
                out.set("logits", logits_full_unpadded)
        return out

    @classmethod
    def _log_probs_generate(cls, tokens, logits, pad_val=-100, pad: bool = True):
        if pad:
            tokens = pad_sequence(
                tokens,
                padding_value=pad_val,
                batch_first=True,
                padding_side="left",
            )
            logits = pad_sequence(
                logits,
                padding_value=0.0,
                batch_first=True,
                padding_side="left",
            )

        # logits = logits.log_softmax(dim=-1)
        # log_probs = logits.gather(-1, tokens.unsqueeze(-1)).squeeze(-1)
        td = TensorDict(logits=logits, tokens=tokens).auto_batch_size_()
        with td.flatten() as tdflat:
            tdflat["log_probs"] = -torch.nn.functional.cross_entropy(
                tdflat["logits"], tdflat["tokens"], reduce=False, ignore_index=pad_val
            )
        td["log_probs"][:, 0] = 0
        log_probs = td["log_probs"]
        return log_probs, logits

    def _compute_log_probs_from_model_output(
        self, model_output, input_ids, attention_mask, pad_val, logits_only=False
    ):
        """Compute log-probs from model output without modifying original tensors.

        Args:
            model_output: Output from the model containing logits
            input_ids: Original input token ids
            attention_mask: Original attention mask
            pad_val: Padding token value to ignore in loss computation
            logits_only: Whether to return only the logits.

        Returns:
            tuple: (log_probs, shifted_logits) where log_probs are the computed log probabilities
                   and shifted_logits are the logits shifted to align with tokens
        """
        logits = model_output["logits"]

        # Create shifted versions for log-prob computation without modifying originals
        shifted_logits = logits[:, :-1, :]
        # shifted_logits = shifted_logits - shifted_logits.logsumexp(dim=-1, keepdim=True)
        shifted_logits = torch.cat(
            [torch.zeros_like(shifted_logits[:, :1]), shifted_logits], 1
        )

        shifted_input_ids = input_ids[:, 1:]
        shifted_input_ids = torch.cat(
            [torch.zeros_like(shifted_input_ids[:, :1]), shifted_input_ids], 1
        )

        # Check that the shape is correct
        if shifted_logits.shape[-2] != shifted_input_ids.shape[-1]:
            raise ValueError(
                f"The logits shape {shifted_logits.shape} does not match the input ids shape {shifted_input_ids.shape}"
            )
        if logits_only:
            return None, shifted_logits

        # Compute log-probs
        td = TensorDict(
            logits=shifted_logits, tokens=shifted_input_ids
        ).auto_batch_size_()
        with td.flatten() as tdflat:
            tdflat["log_probs"] = -torch.nn.functional.cross_entropy(
                tdflat["logits"],
                tdflat["tokens"],
                reduce=False,
                ignore_index=pad_val,
            )
        # For consistency with vllm, we set the log-probs of the first token to 0
        #  However, the first element may not be the first - we want the first of the attention mask,
        #  i.e, the first element that is true on the left
        attention_mask = attention_mask.bool()
        attention_mask_first_left = ~attention_mask[:, :-1] & attention_mask[:, 1:]
        attention_mask_first_left = torch.cat(
            [
                torch.zeros_like(attention_mask_first_left[..., :1]),
                attention_mask_first_left,
            ],
            -1,
        )
        attention_mask_first_left[~(attention_mask_first_left.any(-1)), 0] = True
        assert attention_mask_first_left.any(-1).all()
        attention_mask_first_left = attention_mask_first_left | ~attention_mask
        td["log_probs"][attention_mask_first_left] = 0

        return td["log_probs"], shifted_logits

    def get_dist(
        self,
        tensordict: TensorDictBase,
        tensordict_out: TensorDictBase | None = None,
        logits_key: NestedKey = "logits",
        mask_key: NestedKey | None = None,
        as_padded_tensor: bool | None = None,
        as_nested_tensor: bool | None = None,
        padding_value: float | None = None,
        padding_side: str = "right",
        layout: torch.layout | None = None,
        **kwargs,
    ) -> D.Distribution:
        """Get distribution from logits/log-probs with optional masking.

        This method enables logits computation for distribution creation.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super().get_dist(
                tensordict,
                tensordict_out,
                logits_key,
                mask_key,
                as_padded_tensor,
                as_nested_tensor,
                padding_value,
                padding_side,
                layout,
                **kwargs,
            )
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    def _get_dist_with_prompt_mask(
        self,
        tensordict: TensorDictBase,
        tokens_key: NestedKey = ("tokens", "prompt"),
        logits_key: NestedKey = "logits",
        assistant_mask_key: NestedKey = ("masks", "all_assistant_mask"),
        attention_mask_key: NestedKey = ("masks", "all_attention_mask"),
        **kwargs,
    ) -> D.Distribution:
        """Get distribution masked to only include response tokens (exclude prompt).

        This method enables logits computation for distribution creation.

        This is a provisional method that will be replaced by the `get_dist` method once we have a better masking strategy.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super()._get_dist_with_prompt_mask(
                tensordict,
                tokens_key,
                logits_key,
                assistant_mask_key,
                attention_mask_key,
                **kwargs,
            )
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    def _get_dist_with_assistant_mask(
        self,
        tensordict: TensorDictBase,
        assistant_mask_key: NestedKey = ("masks", "all_assistant_mask"),
        logits_key: NestedKey = "logits",
        **kwargs,
    ) -> D.Distribution:
        """Get distribution masked to only include assistant tokens.

        This method enables logits computation for distribution creation.

        This is a provisional method that will be replaced by the `get_dist` method once we have a better masking strategy.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super()._get_dist_with_assistant_mask(
                tensordict, assistant_mask_key, logits_key, **kwargs
            )
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    def _get_dist_with_attention_mask(
        self,
        tensordict: TensorDictBase,
        attention_mask_key: NestedKey = ("masks", "all_attention_mask"),
        logits_key: NestedKey = "logits",
        **kwargs,
    ) -> D.Distribution:
        """Get distribution masked using attention mask.

        This method enables logits computation for distribution creation.

        This is a provisional method that will be replaced by the `get_dist` method once we have a better masking strategy.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super()._get_dist_with_attention_mask(
                tensordict, attention_mask_key, logits_key, **kwargs
            )
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    def _get_dist_with_custom_mask(
        self,
        tensordict: TensorDictBase,
        mask: torch.Tensor,
        logits_key: NestedKey = "logits",
        **kwargs,
    ) -> D.Distribution:
        """Get distribution with custom mask.

        This method enables logits computation for distribution creation.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super()._get_dist_with_custom_mask(
                tensordict, mask, logits_key, **kwargs
            )
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    # Convenience methods for common LLM training scenarios
    def _get_sft_dist(self, tensordict: TensorDictBase, **kwargs) -> D.Distribution:
        """Get distribution suitable for SFT loss (response tokens only).

        This method enables logits computation for distribution creation.

        This is a provisional method that will be replaced by the `get_dist` method once we have a better masking strategy.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super()._get_sft_dist(tensordict, **kwargs)
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    def _get_rlhf_dist(self, tensordict: TensorDictBase, **kwargs) -> D.Distribution:
        """Get distribution suitable for RLHF loss (assistant tokens only).

        This method enables logits computation for distribution creation.

        This is a provisional method that will be replaced by the `get_dist` method once we have a better masking strategy.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super()._get_rlhf_dist(tensordict, **kwargs)
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    def _get_generic_dist(self, tensordict: TensorDictBase, **kwargs) -> D.Distribution:
        """Get distribution suitable for generic losses (all tokens).

        This method enables logits computation for distribution creation.

        This is a provisional method that will be replaced by the `get_dist` method once we have a better masking strategy.
        """
        self._in_get_dist_call = True
        self.out_keys += ["logits"]
        try:
            return super()._get_generic_dist(tensordict, **kwargs)
        finally:
            self._in_get_dist_call = False
            self.out_keys.remove("logits")

    def _pack_sequences(
        self,
        input_ids: torch.nested.NestedTensor,
        attention_mask: torch.nested.NestedTensor,
    ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
        """Pack sequences into a single tensor."""
        packed_input_ids = input_ids.values()
        lengths = input_ids.lengths()
        if lengths is None:
            offsets = input_ids.offsets()
            lengths = offsets.diff()
            offsets = offsets[1:]
        else:
            offsets = lengths.cumsum(0)
        # Create block-diagonal attention mask to prevent cross-sequence attention
        attention_mask = self._create_block_diagonal_attention_mask(lengths)
        # Create position IDs that restart for each sequence
        position_ids = self._create_packed_position_ids(
            lengths, total_length=packed_input_ids.numel()
        )

        packing_metadata = {
            "sequence_lengths": lengths,
            "cumulative_lengths": offsets,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
        }

        return (
            packed_input_ids.unsqueeze(0),
            attention_mask.unsqueeze(0),
            packing_metadata,
        )

    def _model_forward_with_padded_sequences(
        self,
        tokens_full_padded: torch.Tensor,
        attention_mask_full_padded: torch.Tensor,
        *,
        pad_val: float | int | torch.Tensor | None = None,
        logits_only: bool = False,
        **kwargs,
    ):
        """Forward pass with padded sequences."""
        # Error handling for empty sequences
        if tokens_full_padded.numel() == 0:
            raise ValueError(
                "Input contains empty sequences. Packing/padding requires at least one token per sequence."
            )
        # Error handling for overlong sequences
        config = getattr(self.model, "config", None)
        max_len = getattr(config, "max_position_embeddings", None)
        if max_len is not None and tokens_full_padded.shape[-1] > max_len:
            raise ValueError(
                f"Input sequence length ({tokens_full_padded.shape[-1]}) exceeds model's max_position_embeddings ({max_len}). Consider truncating or splitting your input."
            )
        tokens_out_struct = self.model(
            tokens_full_padded, attention_mask_full_padded, **kwargs
        )
        (
            log_probs_full_padded,
            logits_full_padded,
        ) = self._compute_log_probs_from_model_output(
            tokens_out_struct,
            tokens_full_padded,
            attention_mask_full_padded,
            pad_val,
            logits_only=logits_only,
        )
        return log_probs_full_padded, logits_full_padded

    def _model_forward_with_packed_sequences(
        self,
        flat_input_ids: torch.Tensor,
        block_diag_attention_mask: torch.Tensor,
        *,
        pad: bool = True,
        logits_only: bool = False,
        **kwargs,
    ):
        """Pack sequences into a single tensor and forward them through the model.

        Args:
            flat_input_ids (NestedTensor): NestedTensor of shape (batch_size, -1)
            block_diag_attention_mask (NestedTensor): NestedTensor of shape (batch_size, -1)

        Returns:
            pad (bool): Whether to pad the output tensors.
            logits_only (bool): Whether to return only logits.
            kwargs (dict): Additional keyword arguments to pass to the model.

        """
        # Error handling for empty sequences
        if flat_input_ids.numel() == 0:
            raise ValueError(
                "Input contains empty sequences. Packing requires at least one token per sequence."
            )
        # Error handling for overlong sequences
        # Note: Skipping this check for nested tensors due to symbolic representation issues
        # The model will handle sequence length limits internally
        max_len = getattr(self.model.config, "max_position_embeddings", None)
        if max_len is not None and not hasattr(flat_input_ids, "size"):
            # Only check for regular tensors, not nested tensors
            actual_size = flat_input_ids.shape[-1]
            if actual_size > max_len:
                raise ValueError(
                    f"Input sequence length ({actual_size}) exceeds model's max_position_embeddings ({max_len}). Consider truncating or splitting your input."
                )
        (
            flat_input_ids,
            block_diag_attention_mask,
            packing_metadata,
        ) = self._pack_sequences(flat_input_ids, block_diag_attention_mask)

        outputs = self.model(
            input_ids=flat_input_ids,
            attention_mask=block_diag_attention_mask.unsqueeze(0),
            position_ids=packing_metadata["position_ids"],
            use_cache=False,  # Disable KV cache for packing
            **kwargs,
        )
        log_probs, logits = self._unpack_outputs(
            outputs, packing_metadata, flat_input_ids, pad=pad, logits_only=logits_only
        )
        return log_probs, logits

    def _unpack_outputs(
        self,
        outputs,
        packing_metadata: dict[str, Any],
        flat_input_ids: torch.Tensor,
        pad: bool = True,
        logits_only: bool = False,
    ) -> tuple[torch.Tensor | None, torch.Tensor]:
        """Unpack outputs using nested tensors - zero syncs."""
        # use cross_entropy to compute log_probs
        log_probs, logits = self._compute_log_probs_from_model_output(
            outputs,
            flat_input_ids,
            torch.ones_like(flat_input_ids, dtype=torch.bool),
            -100,
            logits_only=logits_only,
        )
        # check shapes: [1, L] for log_probs, [1, L, vocab_size] for logits
        sequence_lengths = packing_metadata["sequence_lengths"]
        if logits_only:
            log_probs = None
        else:
            if log_probs.shape != logits.shape[:2]:
                raise ValueError(
                    f"Log probs shape {log_probs.shape=} does not match logits shape {logits.shape[:2]=}"
                )
            if log_probs.ndim != 2:
                raise ValueError(f"Log probs shape {log_probs.shape=} is not 2D")
            if logits.ndim != 3:
                raise ValueError(f"Logits shape {logits.shape=} is not 3D")
            if log_probs.shape[1] != sequence_lengths.sum():
                raise ValueError(
                    f"Log probs shape {log_probs.shape=} does not match sequence lengths {sequence_lengths.sum()=}"
                )

            log_probs = log_probs.squeeze(0)
            nested_logprobs = torch.nested.nested_tensor_from_jagged(
                log_probs,
                lengths=sequence_lengths,
            )

        logits = logits.squeeze(0)
        nested_logits = torch.nested.nested_tensor_from_jagged(
            logits,  # Remove batch dim: (total_length, vocab_size)
            lengths=sequence_lengths,
        )

        if logits_only:
            if pad:
                return None, nested_logits.to_padded_tensor(padding=0.0)
            return None, nested_logits
        else:
            if pad:
                return nested_logprobs.to_padded_tensor(
                    padding=0.0
                ), nested_logits.to_padded_tensor(padding=0.0)
            return nested_logprobs, nested_logits

    def _create_block_diagonal_attention_mask(
        self, sequence_lengths: torch.Tensor
    ) -> torch.Tensor:
        """Efficient creation of a block-diagonal attention mask.

        Zero cuda syncs, no integer involved except len(tensor) - compilable.

        Args:
            sequence_lengths: Tensor of shape (batch_size,) containing the lengths of the sequences

        Returns:
            attention_mask: Tensor of shape (batch_size, total_length, total_length)
                where each sequence can only attend to itself.
        """
        seq_ids = torch.arange(len(sequence_lengths), device=sequence_lengths.device)
        position_to_seq_id = seq_ids.repeat_interleave(sequence_lengths)

        attention_mask = position_to_seq_id.unsqueeze(
            1
        ) == position_to_seq_id.unsqueeze(0)
        return attention_mask

    def repeat_interleave_causal(self, sequence_lengths: torch.Tensor) -> torch.Tensor:
        """Same as _create_block_diagonal_attention_mask, but with causal masking."""
        total_length = sequence_lengths.sum()

        seq_ids = torch.arange(len(sequence_lengths), device=sequence_lengths.device)
        position_to_seq_id = seq_ids.repeat_interleave(sequence_lengths)

        positions = torch.arange(int(total_length), device=sequence_lengths.device)

        same_sequence = position_to_seq_id.unsqueeze(1) == position_to_seq_id.unsqueeze(
            0
        )
        causal = positions.unsqueeze(0) <= positions.unsqueeze(1)

        attention_mask = same_sequence & causal
        return attention_mask

    def _create_packed_position_ids(
        self, sequence_lengths: torch.Tensor, total_length: int | None = None
    ) -> torch.Tensor:
        """Create position IDs that restart from 0 for each sequence.

        For sequences of length [3, 2], creates: [0, 1, 2, 0, 1]

        No cuda syncs.
        """
        if total_length is None:
            total_length = int(sequence_lengths.sum().item())

        # Create global position IDs: [0, 1, 2, 3, 4]
        global_positions = torch.arange(total_length, device=sequence_lengths.device)

        # Create sequence start offsets repeated for each position: [0, 0, 0, 3, 3]
        offsets = torch.cat(
            [
                torch.zeros(1, device=sequence_lengths.device),
                sequence_lengths.cumsum(0)[:-1],
            ]
        )
        sequence_starts = offsets.repeat_interleave(sequence_lengths)

        # Subtract to get local positions: [0, 1, 2, 0, 1]
        position_ids = global_positions - sequence_starts

        return position_ids.unsqueeze(0)  # (1, total_length)


class RemoteTransformersWrapper:
    """A remote Ray actor wrapper for TransformersWrapper that provides a simplified interface.

    This class wraps a TransformersWrapper instance as a Ray actor, allowing remote execution
    while providing a clean interface that doesn't require explicit `remote()` and `get()` calls.

    Args:
        model (str): The Hugging Face Transformers model to wrap.
            Must be a string (model name or path) that will be passed to `transformers.AutoModelForCausalLM.from_pretrained`.
            Transformers models are not serializable, so only model names/paths are supported.
        max_concurrency (int, optional): Maximum number of concurrent calls to the remote actor. Defaults to 16.
        validate_model (bool, optional): Whether to validate the model. Defaults to True.
        num_gpus (int, optional): Number of GPUs to use. Defaults to 0.
        num_cpus (int, optional): Number of CPUs to use. Defaults to 0.
        **kwargs: All other arguments are passed directly to TransformersWrapper.

    Example:
        >>> import ray
        >>> from torchrl.modules.llm.policies import RemoteTransformersWrapper
        >>>
        >>> # Initialize Ray if not already done
        >>> if not ray.is_initialized():
        ...     ray.init()
        >>>
        >>> # Create remote wrapper
        >>> remote_wrapper = RemoteTransformersWrapper(
        ...     model="gpt2",
        ...     input_mode="history",
        ...     generate=True,
        ...     generate_kwargs={"max_new_tokens": 50}
        ... )
        >>>
        >>> # Use like a regular wrapper (no remote/get calls needed)
        >>> result = remote_wrapper(tensordict_input)
        >>> print(result["text"].response)
    """

    def __init__(
        self,
        model,
        max_concurrency: int = 16,
        validate_model: bool = True,
        actor_name: str | None = None,
        num_gpus: int = 1,
        num_cpus: int = 1,
        **kwargs,
    ):
        import ray

        # Validate model parameter - only strings are allowed for Transformers
        if not isinstance(model, str) and validate_model:
            raise ValueError(
                "For RemoteTransformersWrapper, the model parameter must be a string "
                f"(model name or path). Got type: {type(model)}. "
                "Transformers models are not serializable, so only model names/paths are supported. "
                "You can bypass this check by setting validate_model=False."
            )

        if not ray.is_initialized():
            ray.init()

        if actor_name is not None:
            # Check if an actor with this name already exists
            try:
                existing_actor = ray.get_actor(actor_name)
                # If we can get the actor, assume it's alive and use it
                self._remote_wrapper = existing_actor
                torchrl_logger.info(f"Using existing actor {actor_name}")
                return
            except ValueError:
                # Actor doesn't exist, create a new one
                torchrl_logger.info(f"Creating new actor {actor_name}")

        # Create the remote actor with the unique name
        self._remote_wrapper = (
            ray.remote(TransformersWrapper)
            .options(
                max_concurrency=max_concurrency,
                name=actor_name,
                num_gpus=num_gpus,
                num_cpus=num_cpus,
            )
            .remote(model, **kwargs)
        )

    def __call__(self, tensordict, **kwargs):
        """Forward pass that automatically handles remote execution."""
        import ray

        return ray.get(self._remote_wrapper.forward.remote(tensordict, **kwargs))

    def get_new_version(self, **kwargs):
        """Get a new version of the wrapper with altered parameters."""
        import ray

        return ray.get(self._remote_wrapper.get_new_version.remote(**kwargs))

    def get_dist(self, tensordict, **kwargs):
        """Get distribution from logits/log-probs with optional masking."""
        import ray

        return ray.get(self._remote_wrapper.get_dist.remote(tensordict, **kwargs))

    def get_dist_with_prompt_mask(self, tensordict, **kwargs):
        """Get distribution masked to only include response tokens (exclude prompt)."""
        import ray

        return ray.get(
            self._remote_wrapper.get_dist_with_prompt_mask.remote(tensordict, **kwargs)
        )

    def _get_dist_with_assistant_mask(self, tensordict, **kwargs):
        """Get distribution masked to only include assistant tokens."""
        import ray

        return ray.get(
            self._remote_wrapper._get_dist_with_assistant_mask.remote(
                tensordict, **kwargs
            )
        )

    def _get_dist_with_attention_mask(self, tensordict, **kwargs):
        """Get distribution masked using attention mask."""
        import ray

        return ray.get(
            self._remote_wrapper._get_dist_with_attention_mask.remote(
                tensordict, **kwargs
            )
        )

    def _get_dist_with_custom_mask(self, tensordict, **kwargs):
        """Get distribution with custom mask."""
        import ray

        return ray.get(
            self._remote_wrapper._get_dist_with_custom_mask.remote(tensordict, **kwargs)
        )

    def _get_sft_dist(self, tensordict, **kwargs):
        """Get distribution suitable for SFT loss (response tokens only)."""
        import ray

        return ray.get(self._remote_wrapper._get_sft_dist.remote(tensordict, **kwargs))

    def _get_rlhf_dist(self, tensordict, **kwargs):
        """Get distribution suitable for RLHF loss (assistant tokens only)."""
        import ray

        return ray.get(self._remote_wrapper._get_rlhf_dist.remote(tensordict, **kwargs))

    def _get_generic_dist(self, tensordict, **kwargs):
        """Get distribution suitable for generic losses (all tokens)."""
        import ray

        return ray.get(
            self._remote_wrapper._get_generic_dist.remote(tensordict, **kwargs)
        )

    def log_prob(self, data, **kwargs):
        """Compute log probabilities."""
        import ray

        return ray.get(self._remote_wrapper.log_prob.remote(data, **kwargs))

    def cleanup_batching(self):
        """Clean up batching resources."""
        import ray

        return ray.get(self._remote_wrapper.cleanup_batching.remote())

    def __del__(self):
        """Cleanup when the wrapper is destroyed."""
        try:
            import ray

            if hasattr(self, "_remote_wrapper") and ray.is_initialized():
                # Clean up batching resources
                try:
                    ray.get(self._remote_wrapper.cleanup_batching.remote())
                except Exception:
                    pass  # Ignore cleanup errors during destruction
        except Exception:
            pass  # Ignore any errors during cleanup

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit with cleanup."""
        self.cleanup_batching()

    def get_batching_state(self):
        """Get the current batching state."""
        import ray

        return ray.get(self._remote_wrapper.get_batching_state.remote())

    @property
    def generate(self):
        """Whether text generation is enabled."""
        import ray

        return ray.get(self._remote_wrapper.generate.remote)

    @property
    def pad_output(self):
        """Whether output sequences are padded."""
        import ray

        return ray.get(self._remote_wrapper.pad_output.remote)

    @property
    def text_key(self):
        """The key for text output."""
        import ray

        return ray.get(self._remote_wrapper.text_key.remote)

    @property
    def tokens_key(self):
        """The key for tokens output."""
        import ray

        return ray.get(self._remote_wrapper.tokens_key.remote)

    @property
    def masks_key(self):
        """The key for masks output."""
        import ray

        return ray.get(self._remote_wrapper.masks_key.remote)

    @property
    def log_probs_key(self):
        """The key for log probabilities output."""
        import ray

        return ray.get(self._remote_wrapper.log_probs_key.remote)

    @property
    def in_keys(self):
        """The input keys."""
        import ray

        return ray.get(self._remote_wrapper.in_keys.remote)

    @property
    def out_keys(self):
        """The output keys."""
        import ray

        return ray.get(self._remote_wrapper.out_keys.remote)

    @property
    def inplace(self):
        """Whether in-place operations are used."""
        import ray

        return ray.get(self._remote_wrapper.inplace.remote)

    @property
    def device(self):
        """The device used for computation."""
        import ray

        return ray.get(self._remote_wrapper.device.remote)

    @property
    def layout(self):
        """The layout used for output tensors."""
        import ray

        return ray.get(self._remote_wrapper.layout.remote)

    @property
    def num_samples(self):
        """The number of samples to generate."""
        import ray

        return ray.get(self._remote_wrapper.num_samples.remote)

    @property
    def batching(self):
        """Whether batching is enabled."""
        import ray

        return ray.get(self._remote_wrapper.batching.remote)

    @property
    def collector(self):
        """The collector associated with the module."""
        import ray

        return ray.get(self._remote_wrapper.collector.remote)

    @property
    def log_prob_keys(self):
        """The keys for log probabilities."""
        import ray

        return ray.get(self._remote_wrapper.log_prob_keys.remote)

    @log_prob_keys.setter
    def log_prob_keys(self, value):
        """Set the keys for log probabilities."""
        import ray

        ray.get(self._remote_wrapper.log_prob_keys.remote(value))

    @property
    def dist_params_keys(self):
        """The keys for distribution parameters."""
        import ray

        return ray.get(self._remote_wrapper.dist_params_keys.remote)

    @property
    def dist_sample_keys(self):
        """The keys for distribution samples."""
        import ray

        return ray.get(self._remote_wrapper.dist_sample_keys.remote)
